import asyncio
import cProfile
from datetime import datetime
import time
import logging
from logging.handlers import RotatingFileHandler
import os
import re
import random
import shutil
import threading
from wxauto import WeChat
from openai import OpenAI

# 配置参数集中管理
class Config:
    LISTEN_LIST = [['AI测试', '提示词示例3'],['深大-AI测试3', '提示词示例1']]

    DEEPSEEK_API_KEY = '330f54b7-915b-48a9-9599-c5e47c118cb5'
    DEEPSEEK_BASE_URL = 'https://ark.cn-beijing.volces.com/api/v3'
    MODEL = 'doubao-1-5-pro-256k-250115'
    MAX_TOKEN = 2000
    TEMPERATURE = 1.1
    QUEUE_WAITING_TIME = 7
    AUTO_MESSAGE = '请你模拟系统设置的角色，在微信上找对方继续刚刚的话题或者询问对方在做什么'
    ENABLE_AUTO_MESSAGE = False
    MIN_COUNTDOWN_HOURS = 1.0
    MAX_COUNTDOWN_HOURS = 2.0
    QUIET_TIME_START = '22:00'
    QUIET_TIME_END = '8:00'
    AVERAGE_TYPING_SPEED = 0.01
    RANDOM_TYPING_SPEED_MIN = 0.01
    RANDOM_TYPING_SPEED_MAX = 0.01
    ENABLE_MEMORY = True
    MEMORY_TEMP_DIR = 'Memory_Temp'
    MAX_MESSAGE_LOG_ENTRIES = 4
    MAX_MEMORY_NUMBER = 10
    Accept_All_Group_Chat_Messages = False

# 日志配置优化
log_file = 'bot.log'
max_bytes = 10 * 1024 * 1024  # 10MB
backup_count = 5
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        RotatingFileHandler(log_file, maxBytes=max_bytes, backupCount=backup_count),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 单例模式管理微信实例
class WeChatManager:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        with cls._lock:
            if not cls._instance:
                cls._instance = super().__new__(cls)
                self = cls._instance
                self.wx = WeChat()
                self.ROBOT_WX_NAME = self.wx.A_MyIcon.Name
                self.initialize_listen()
        return cls._instance

    def initialize_listen(self):
        for user in Config.LISTEN_LIST:
            try:
                self.wx.AddListenChat(who=user[0], savepic=True)
            except Exception as e:
                logger.error(f"添加监听对象 {user[0]} 时出错: {e}")

# 用户状态管理
class UserState:
    def __init__(self):
        self.queues = {}
        self.contexts = {}
        self.timers = {}
        self.wait_times = {}
        self.lock = threading.RLock()

    def get_queue(self, user_id):
        with self.lock:
            return self.queues.setdefault(user_id, {
                'messages': [],
                'sender_name': '',
                'username': '',
                'last_message_time': 0,
                'processed_messages': set()
            })

    def update_timer(self, user):
        with self.lock:
            self.timers[user] = time.time()
            self.wait_times[user] = random.uniform(
                Config.MIN_COUNTDOWN_HOURS,
                Config.MAX_COUNTDOWN_HOURS
            ) * 3600

# 核心功能模块
class ChatBot:
    def __init__(self):
        self.wx_manager = WeChatManager()
        self.user_state = UserState()
        self.quiet_start = datetime.strptime(Config.QUIET_TIME_START, "%H:%M").time()
        self.quiet_end = datetime.strptime(Config.QUIET_TIME_END, "%H:%M").time()
        self.root_dir = os.path.dirname(os.path.abspath(__file__))
        self.memory_temp_dir = os.path.join(self.root_dir, Config.MEMORY_TEMP_DIR)
        os.makedirs(self.memory_temp_dir, exist_ok=True)
        self.client = OpenAI(
            api_key=Config.DEEPSEEK_API_KEY,
            base_url=Config.DEEPSEEK_BASE_URL
        )
        # zjq,消息监听时延
        self.wait = 0.1
        self.prompt_cache = {}

    def is_quiet_time(self):
        now = datetime.now().time()
        if self.quiet_start <= self.quiet_end:
            return self.quiet_start <= now <= self.quiet_end
        return now >= self.quiet_start or now <= self.quiet_end

    async def process_message(self, msg, who):
        try:
            username = who
            content = getattr(msg, 'content', None) or getattr(msg, 'text', None)

            if not content:
                return

            logger.info(f'"接受消息"【{username}】：{content}')
            # zjq
            self.log_memory(who, username+":"+content)

            self.user_state.update_timer(username)
            await self.handle_message(username, content)

        except Exception as e:
            logger.error(f"消息处理失败: {str(e)}", exc_info=True)

    async def handle_message(self, user_id, content):
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        content = f"[{current_time}] {content}"

        with self.user_state.lock:
            queue = self.user_state.get_queue(user_id)
            if content in queue['processed_messages']:
                return
            queue['sender_name'] = user_id
            queue['username'] = user_id
            queue['messages'].append(content)
            queue['last_message_time'] = time.time()
            queue['processed_messages'].add(content)

            if len(queue['messages']) > 5:
                queue['messages'].pop(0)

        if not self.is_sending_message():
            await self.process_user_messages(user_id)

    def is_sending_message(self):
        return False  # 需根据实际情况实现

    async def process_user_messages(self, user_id):
        # zjq queue
        queue = self.user_state.get_queue(user_id)
        messages = queue['messages']
        # merged = ' '.join(messages)


        # messages = queue['messages']
        # if len(messages) > 5:
        #     messages = messages[-5:]
        # if messages:
        #     for i in range(len(messages)):
        #         if i == 0:
        #             messages[i] = f"最新消息:{messages[i]}"
        #         else:
        #             messages[i] = f"历史对话:{messages[i]}"
        #     merged = "\n".join(messages)
        # else:
        #     merged = None


        if messages:  # 先判断列表是否为空，避免空列表取值报错
            merged = messages[-1]
        else:
            merged = None
        logger.info(f"调用llm开始,消息：{merged}")
        reply = await self.generate_response(merged, user_id)
        logger.info(f"调用llm结束")

        await self.send_reply(user_id, queue['sender_name'], merged, reply)

    async def generate_response(self, message, user_id):
        try:
            user_prompt = self.get_user_prompt(user_id)
            response = self.client.chat.completions.create(
                model=Config.MODEL,
                messages=[
                    {"role": "system", "content": user_prompt},
                    {"role": "user", "content": message}
                ],
                temperature=Config.TEMPERATURE,
                max_tokens=Config.MAX_TOKEN,
                stream=False
            )
            logger.info(f"调用llm开始,user_prompt {user_prompt} message {message} ")

            return response.choices[0].message.content.strip()
        except Exception as e:
            logger.error(f"生成回复失败: {str(e)}", exc_info=True)
            return "抱歉，我现在有点忙，稍后再聊吧。"

    async def send_reply(self, user_id, sender_name, message, reply):
        logger.info(f"回复开始")

        try:
            reply = self.remove_timestamps(reply)
            parts = [p.strip() for p in reply.split('\\') if p.strip()]
            for i, part in enumerate(parts):
                self.wx_manager.wx.SendMsg(part, user_id)
                logger.debug(f"回复完成 {sender_name}: {part}")
                self.log_memory(user_id, "回复："+part)

                if i < len(parts) - 1:
                    delay = len(parts[i + 1]) * (
                            Config.AVERAGE_TYPING_SPEED +
                            random.uniform(Config.RANDOM_TYPING_SPEED_MIN, Config.RANDOM_TYPING_SPEED_MAX)
                    )
                    # await asyncio.sleep(max(delay, 1))
        except Exception as e:
            logger.error(f"发送回复失败: {str(e)}", exc_info=True)

    def log_memory(self, user_id, content):
        if not Config.ENABLE_MEMORY:
            return

        prompt_name = next((u[1] for u in Config.LISTEN_LIST if u[0] == user_id), user_id)
        log_file = os.path.join(self.memory_temp_dir, f'{user_id}_{prompt_name}_log.txt')

        if os.path.exists(log_file) and os.path.getsize(log_file) > 1 * 1024 * 1024:
            archive = os.path.join(self.memory_temp_dir, f'{user_id}_log_archive_{int(time.time())}.txt')
            shutil.move(log_file, archive)

        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} | {content}\n")

    def get_user_prompt(self, user_id):
        prompt_name = next((u[1] for u in Config.LISTEN_LIST if u[0] == user_id), user_id)
        # zjq每次都重新读取文件，不使用缓存

        # if prompt_name in self.prompt_cache:
        #     return self.prompt_cache[prompt_name]

        prompt_path = os.path.join(self.root_dir, 'prompts', f'{prompt_name}.md')

        if not os.path.exists(prompt_path):
            raise FileNotFoundError(f"Prompt文件 {prompt_name}.md 未找到")

        with open(prompt_path, 'r', encoding='utf-8') as f:
            content = f.read()
            print(content)
            self.prompt_cache[prompt_name] = content
            return content

    def remove_timestamps(self, text):
        timestamp_pattern = r'\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\]'
        return re.sub(timestamp_pattern, '', text).strip()

    async def start_listening(self):
        logger.info(f'"start_listening"')

        while True:
            try:
                if self.wx_manager.wx is None:
                    self.wx_manager = WeChatManager()
                    await asyncio.sleep(0.1)
                    logger.info(f'"创建 WeChatManager"')

                if self.wx_manager.wx is not None:
                    logger.debug(f'"wx_manager.wx.GetListenMessage——start"')

                    msgs = self.wx_manager.wx.GetListenMessage()
                    logger.debug(f'"wx_manager.wx.GetListenMessage-end"')

                    for chat in msgs:
                        who = chat.who
                        one_msgs = msgs.get(chat)
                        logger.debug(f'"f" for chat in msgs {one_msgs}.md 未找到""')

                        for msg in one_msgs:
                            msgtype = msg.type
                            content = msg.content
                            if not content:
                                continue
                            if msgtype != 'friend':
                                logger.debug(f"非好友消息，忽略! 消息类型: {msgtype}")
                                continue
                            if who == msg.sender:
                                await self.process_message(msg, who)
                            elif Config.Accept_All_Group_Chat_Messages:
                                if not msg.content.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                                    msg.content = f"群聊消息[{msg.sender}]:" + msg.content
                                await self.process_message(msg, who)
                            elif self.wx_manager.ROBOT_WX_NAME and (
                                    re.search(f'@{self.wx_manager.ROBOT_WX_NAME}\u2005', msg.content)):
                                msg.content = re.sub(f'@{self.wx_manager.ROBOT_WX_NAME}\u2005', '', content).strip()
                                msg.content = f"群聊消息[{msg.sender}]:" + msg.content
                                await self.process_message(msg, who)
                            elif (bool(re.search("邀请", msg.content))):
                                # 处理群聊信息，只有@当前机器人才会处理
                                msg.content = re.sub(f'@{self.wx_manager.ROBOT_WX_NAME}\u2005', '', content).strip()
                                msg.content = "群聊消息[" + msg.sender + "]:" + msg.content
                                logger.info(f'"process_message——start"')
                                await self.process_message(msg, who)
                                logger.info(f'"process_message-end"')

                            elif msgtype == 'group' and '@' in msg.content:
                                msg.content = f"群聊消息[{msg.sender}]:" + msg.content
                                await self.process_message(msg, who)
                            else:
                                logger.debug(f"非需要处理消息: {content}")
                else:
                    logger.error("微信实例初始化失败，等待重试...")
                    await asyncio.sleep(0.1)
            except Exception as e:
                logger.error(f"监听错误: {str(e)}", exc_info=True)
                print(
                    "\033[31m重要提示：请不要关闭程序打开的微信聊天框！若命令窗口收不到消息，请将微信聊天框置于最前台！ \033[0m")
                self.wx_manager.wx = None
                await asyncio.sleep(0.1)
            await asyncio.sleep(self.wait)

    async def schedule_tasks(self):
        while True:
            try:
                if Config.ENABLE_AUTO_MESSAGE:
                    await self.check_user_timeouts()  # 添加await
                if Config.ENABLE_MEMORY:
                    await self.memory_manager()  # 添加await
            except Exception as e:
                logger.error(f"定时任务执行出错: {str(e)}", exc_info=True)
            await asyncio.sleep(1)

    async def check_user_timeouts(self):  # 添加async
        current_time = time.time()
        with self.user_state.lock:
            for user in list(self.user_state.timers.keys()):
                last_active = self.user_state.timers[user]
                wait_time = self.user_state.wait_times[user]
                if current_time - last_active >= wait_time:
                    if not self.is_quiet_time():
                        current_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        auto_content = f"[{current_time_str}] {Config.AUTO_MESSAGE}"
                        logger.info(f"为用户 {user} 发送自动消息:{auto_content}")
                        reply = await self.generate_response(auto_content, user)  # 使用await
                        await self.send_reply(user, user, auto_content, reply)  # 使用await
                    self.user_state.update_timer(user)

    async def memory_manager(self):  # 添加async
        for user in [entry[0] for entry in Config.LISTEN_LIST]:
            prompt_name = next((u[1] for u in Config.LISTEN_LIST if u[0] == user), user)
            log_file = os.path.join(self.memory_temp_dir, f'{user}_{prompt_name}_log.txt')
            try:
                if os.path.exists(log_file):
                    with open(log_file, 'r', encoding='utf-8') as f:
                        line_count = sum(1 for _ in f)
                    if line_count >= Config.MAX_MESSAGE_LOG_ENTRIES:
                        await self.summarize_and_save(user)  # 添加await
            except Exception as e:
                logger.error(f"记忆管理失败: {str(e)}", exc_info=True)

    async def summarize_and_save(self, user_id):  # 添加async
        prompt_name = next((u[1] for u in Config.LISTEN_LIST if u[0] == user_id), user_id)
        log_file = os.path.join(self.memory_temp_dir, f'{user_id}_{prompt_name}_log.txt')
        if not os.path.exists(log_file):
            logger.warning(f"日志文件不存在: {log_file}")
            return
        if os.path.getsize(log_file) == 0:
            logger.info(f"空日志文件: {log_file}")
            return

        with open(log_file, 'r', encoding='utf-8') as f:
            logs = [line.strip() for line in f if line.strip()]
            if len(logs) < Config.MAX_MESSAGE_LOG_ENTRIES:
                logger.info(f"日志条目不足（{len(logs)}条），无需处理")
                return

        full_logs = '\n'.join(logs)
        summary_prompt = f"请用中文总结以下对话，提取重要信息形成记忆片段的摘要（仅输出摘要不要输出其它信息）：\n{full_logs}"
        logger.info(f"summary_prompt ({summary_prompt})")
        summary = await self.generate_response(summary_prompt, "system")  # 使用await

        summary = re.sub(
            r'\*{0,2}(重要度|摘要)\*{0,2}[\s:]*\d*[\.]?\d*[\s\\]*|## 记忆片段 \[\d{4}-\d{2}-\d{2} \d{2}:\d{2}\]',
            '',
            summary,
            flags=re.MULTILINE
        ).strip()

        importance_prompt = f"为以下内容的重要性评分（1 - 5，直接回复数字）：\n{summary}"
        importance_response = await self.generate_response(importance_prompt, "system")  # 使用await
        importance_match = re.search(r'[1-5]', importance_response)
        if importance_match:
            importance = min(max(int(importance_match.group()), 1), 5)
        else:
            importance = 3
            logger.warning(f"无法解析重要性评分，使用默认值3。原始响应：{importance_response}")

        current_time = datetime.now().strftime("%Y-%m-%d %H:%M")
        memory_entry = f"""## 记忆片段 [{current_time}]
**重要度**: {importance}
**摘要**: {summary}

"""

        prompts_dir = os.path.join(self.root_dir, 'prompts')
        os.makedirs(prompts_dir, exist_ok=True)
        user_prompt_file = os.path.join(prompts_dir, f'{prompt_name}.md')
        temp_file = f"{user_prompt_file}.tmp"
        backup_file = f"{user_prompt_file}.bak"

        try:
            with open(temp_file, 'w', encoding='utf-8') as f:
                if os.path.exists(user_prompt_file):
                    with open(user_prompt_file, 'r', encoding='utf-8') as src:
                        f.write(src.read().rstrip() + '\n\n')
                f.write(memory_entry)

            if os.path.exists(user_prompt_file):
                shutil.copyfile(user_prompt_file, backup_file)
            shutil.move(temp_file, user_prompt_file)
        except Exception as e:
            logger.error(f"保存记忆片段失败: {str(e)}", exc_info=True)
            if os.path.exists(backup_file):
                shutil.move(backup_file, user_prompt_file)
        finally:
            with open(log_file, 'w', encoding='utf-8') as f:
                f.truncate()

# 初始化和入口
async def main():
    try:
        bot = ChatBot()

        # 启动监听任务
        asyncio.create_task(bot.start_listening())

        # 启动定时任务
        asyncio.create_task(bot.schedule_tasks())

        while True:
            await asyncio.sleep(0.1)

    except KeyboardInterrupt:
        logger.info("用户终止程序")
    except Exception as e:
        logger.error(f"发生异常: {str(e)}", exc_info=True)

if __name__ == '__main__':
    asyncio.run(main())
    # cProfile.run('main()')
